from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser
from openai import OpenAI as ChatGptOpenAI

from common import TranslateLanguage

llm = ChatOpenAI(model_name='gpt-3.5-turbo', max_tokens=256, temperature=0)

client = ChatGptOpenAI()


class OpenaiModel:
    @classmethod
    def getTransPromptMessage(cls):
        system_template = """
        你是个翻译专家，精通多国语言，将下面的内容从{fromLanguage}翻译成{toLanguage}
        """
        system_prompt = SystemMessagePromptTemplate.from_template(system_template)
        human_prompt = HumanMessagePromptTemplate.from_template("{content}")

        messages = [system_prompt, human_prompt]
        print(messages)
        return messages

    def request_openai(self, content: str, fromLanguage: str, toLanguage: str):
        f_language = TranslateLanguage.getLanguage(fromLanguage)
        t_language = TranslateLanguage.getLanguage(toLanguage)
        messages = self.getTransPromptMessage()
        prompt = ChatPromptTemplate.from_messages(messages)
        chain = prompt | llm | StrOutputParser()
        print(chain)
        return chain.invoke({"fromLanguage": f_language, "toLanguage": t_language, "content": content})
